'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: backbone_zoo.py
@time: 2020-05-29 11:17:12
@desc: 
'''
from ELib.utils.register import Registry
import torch

BACKBONE_ZOO = Registry("BACKBONE")


def get_segmentation_backbone(cfg, backbone, norm_layer=torch.nn.BatchNorm2d):
    """
    Built the backbone model, defined by `cfg.MODEL.BACKBONE`.
    """
    model = BACKBONE_ZOO.get(backbone)(cfg, norm_layer)
    return model


class BackboneSeg(torch.nn.Module):
    def __init__(self, cfg):
        super(BackboneSeg, self).__init__()
        self.cfg = cfg

    def load_weights(self, weights):
        self.load_state_dict(weights, strict=False)
        return None

